#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2019/8/20
# @Author  : github.com/guofei9987


import numpy as np
import tqdm

from .GA import GeneticAlgorithmBase


class DE(GeneticAlgorithmBase):
    def __init__(self, func, n_dim, F=0.5,
                 size_pop=50, max_iter=200, prob_mut=0.3, max_stay_iter=None,
                 lb=-1, ub=1,
                 constraint_eq=tuple(), constraint_ueq=tuple()):
        super().__init__(func, n_dim, size_pop, max_iter, prob_mut,
                         constraint_eq=constraint_eq, constraint_ueq=constraint_ueq)

        self.F = F
        self.V, self.U = None, None
        self.max_stay_iter = max_stay_iter or max_iter
        self.lb, self.ub = np.array(lb) * np.ones(self.n_dim), np.array(ub) * np.ones(self.n_dim)
        self.crtbp()

    def crtbp(self):
        # create the population
        self.X = np.random.uniform(low=self.lb, high=self.ub, size=(self.size_pop, self.n_dim)).astype(np.float32)
        return self.X

    def chrom2x(self, Chrom):
        pass

    def ranking(self):
        pass

    def mutation(self):
        '''
        V[i]=X[r1]+F(X[r2]-X[r3]),
        where i, r1, r2, r3 are randomly generated
        '''
        X = self.X
        # i is not needed,
        # and TODO: r1, r2, r3 should not be equal
        random_idx = np.random.randint(0, self.size_pop, size=(self.size_pop, 3))

        r1, r2, r3 = random_idx[:, 0], random_idx[:, 1], random_idx[:, 2]

        # 这里F用固定值，为了防止早熟，可以换成自适应值
        self.V = X[r1, :] + self.F * (X[r2, :] - X[r3, :])

        # the lower & upper bound still works in mutation
        mask = np.random.uniform(low=self.lb, high=self.ub, size=(self.size_pop, self.n_dim)).astype(np.float32)
        self.V = np.where(self.V < self.lb, mask, self.V)
        self.V = np.where(self.V > self.ub, mask, self.V)
        return self.V

    def crossover(self):
        '''
        if rand < prob_crossover, use V, else use X
        '''
        mask = np.random.rand(self.size_pop, self.n_dim) < self.prob_mut
        self.U = np.where(mask, self.V, self.X)
        return self.U

    def selection(self):
        '''
        greedy selection
        '''
        X = self.X.copy()

        self.X = np.concatenate((self.X, self.U), axis=0)
        f_U = self.x2y()

        self.X = self.X[f_U.flatten() <= np.median(f_U)][:X.shape[0]]

        return self.X

    def run(self, max_iter=None, max_stay_iter=None):
        best_step = 0
        self.max_stay_iter = max_stay_iter or self.max_stay_iter
        self.max_iter = max_iter or self.max_iter

        generation_best_X = np.zeros_like([self.X[0]]).repeat(self.max_iter, axis=0)

        all_history_Y, all_history_FitV, generation_best_Y = \
            np.zeros([self.max_iter, 1], dtype=np.float32), \
            np.zeros([self.max_iter, 1], dtype=np.float32), \
            np.zeros([self.max_iter, 1], dtype=np.float32)

        with tqdm.tqdm(range(self.max_iter)) as pbar:
            for i in pbar:
                self.mutation()
                self.crossover()
                self.selection()
                self.Y = self.x2y()

                # record the best ones
                generation_best_index = self.Y.argmin()
                generation_best_X[i], generation_best_Y[i] = \
                    self.X[generation_best_index], \
                    self.Y[generation_best_index]

                all_history_Y[i] = self.Y[generation_best_index]

                if i == 0 or generation_best_Y[i] < np.min(generation_best_Y[:i]):
                    pbar.set_description('%g@%d' % (generation_best_Y[i], i))
                    best_step = i
                if i - best_step > self.max_stay_iter:
                    break
                cnt = i + 1

        self.generation_best_X += list(generation_best_X[:cnt])
        self.generation_best_Y += list(generation_best_Y[:cnt])
        self.all_history_Y += list(all_history_Y[:cnt])

        global_best_index = np.array(self.generation_best_Y).argmin()
        self.best_x = self.generation_best_X[global_best_index]
        self.best_y = self.func(np.array([self.best_x]))
        return self.best_x, self.best_y
